package org.hhf.ai.service;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.hhf.ai.config.AiConfig;
import org.hhf.ai.constant.AIConstant;
import org.hhf.ai.pojo.AiChatForm;
import org.hhf.ai.pojo.AiChatVo;
import org.hhf.ai.tools.SSEUtils;
import org.hhf.ai.tools.funcalling.FunctionRegistry;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
@Service
public class AiChatService {
    @Autowired
    private WebClient webClient;
    @Autowired
    private AiConfig aiConfig;

    public SseEmitter sseChat(AiChatForm form) {
        Long conversationId = form.getConversationId();
        setHistory(conversationId, AIConstant.ROLE_ASSISTANT, form.getPrompt());
        String sessionId = form.getSessionId();
        SseEmitter emitter = SSEUtils.create(sessionId);
        // 构建请求体
        Map<String, Object> requestBody = new HashMap<>();
        buildBody(form,requestBody);
        AtomicInteger totalTokens = new AtomicInteger();
        //定义字符串记录流式响应结果
        StringBuilder responseBuilder = new StringBuilder();
        StringBuilder reasoningBuilder = new StringBuilder();
        StringBuilder functionNameBuilder = new StringBuilder();
        StringBuilder argumentsBuilder = new StringBuilder();
        CompletableFuture.runAsync(() -> {
            try {
                webClient.post()
                        .uri("/v1/chat/completions")
                        .bodyValue(requestBody)
                        .accept(MediaType.TEXT_EVENT_STREAM) // 接收流式响应
                        .retrieve()
                        .bodyToFlux(String.class)
                        .subscribe(response -> {
                                    // 解析流式响应的 JSON 数据
                                    try {

                                        if ("[DONE]".equals(response)) {
                                            SSEUtils.send(sessionId, new AiChatVo(AIConstant.TYPE_TEXT, response));
                                        } else {
                                            JsonNode jsonNode = new ObjectMapper().readTree(response);

                                            if (jsonNode.has("choices")) {
                                                JsonNode choices = jsonNode.get("choices");
                                                if (choices.isArray() && choices.size() > 0) {
                                                    JsonNode subChoice = choices.get(0);
                                                    if (subChoice.has("finish_reason")) {
                                                        String finishReason = subChoice.get("finish_reason").asText();
                                                        if ("stop".equals(finishReason)) {
                                                            JsonNode usage = subChoice.get("usage");
                                                            totalTokens.set(getTotalTokens(usage));
                                                            SSEUtils.send(sessionId, new AiChatVo(AIConstant.TYPE_TEXT,""));
                                                        }
                                                    }
                                                    JsonNode delta = subChoice.get("delta");
                                                    if (delta != null) {
                                                        JsonNode content = delta.get("content");
                                                        JsonNode reasoning_content = delta.get("reasoning_content");
                                                        JsonNode toolCalls = delta.get("tool_calls");

                                                        if (reasoning_content != null) {
                                                            String ctx = reasoning_content.asText();
                                                            if (StringUtils.isNotBlank(ctx) && !"null".equals(ctx)) {
                                                                reasoningBuilder.append(ctx);
                                                                SSEUtils.send(sessionId, new AiChatVo(AIConstant.TYPE_THINKING, ctx));
                                                            }
                                                        }
                                                        if (content != null) {
                                                            String ctx = content.asText();
                                                            if (StringUtils.isNotBlank(ctx) && !"null".equals(ctx)) {
                                                                responseBuilder.append(ctx);
                                                                SSEUtils.send(sessionId, new AiChatVo(AIConstant.TYPE_TEXT, ctx));
                                                            }
                                                        }
                                                        if (toolCalls != null) {
                                                            JsonNode function = toolCalls.get(0).get("function");
                                                            JsonNode functionNameNode = function.get("name");

                                                            if (functionNameNode != null) {
                                                                String name = functionNameNode.asText();
                                                                if (StringUtils.isNotBlank(name) && !"null".equals(name)) {
                                                                    functionNameBuilder.append(name);
                                                                    log.info("--------------------------name:{}", name);
                                                                }
                                                            }
                                                            JsonNode argumentsNode = function.get("arguments");
                                                            if (argumentsNode != null) {
                                                                String arguments = argumentsNode.asText();
                                                                if (StringUtils.isNotBlank(arguments) && !"null".equals(arguments)) {
                                                                    argumentsBuilder.append(arguments);
                                                                    log.info("--------------------------arguments:{}", arguments);
                                                                }
                                                            }

                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    } catch (Exception e) {
                                        log.error("数据解析失败: " + e);
//                                    emitter.completeWithError(e);
                                        SSEUtils.close(sessionId);
                                    }
                                },
                                error -> {
                                    log.error("subscribe error:{}", error);
                                    SSEUtils.close(sessionId);
//                            emitter.completeWithError(error);
                                },
                                ()-> {
                                    // 如果需要执行tools函数，则递归执行，否则emitter.complete()关闭客户端SSE连接
                                    if (functionNameBuilder.length() > 0) {
                                        try {
                                            JsonNode argumentsNode = new ObjectMapper().readTree(argumentsBuilder.toString());
                                            String functionName = functionNameBuilder.toString();
                                            String functionResult = FunctionRegistry.getInstance().functionInvoker(functionName, argumentsNode);
                                            String prompt = form.getPrompt()+"请分析以下数据并给出结果：" + functionResult;
                                            form.setPrompt(prompt);
                                            log.info("function call 生成的提示词：{}", prompt);
                                            // 递归调用 streamChatSse，并等待完成
                                            sseChat(form);
                                        } catch (JsonProcessingException e) {
                                            log.error(e.getMessage(), e);
                                        }
                                    } else {
                                        log.info("Stream completed.");
                                        SSEUtils.close(sessionId);
                                    }
                                });

            } catch (Exception e) {
                log.error(e.getMessage(), e);
//                emitter.completeWithError(e);
                SSEUtils.close(sessionId);
            }
        });

        // 连接状态监听
        emitter.onCompletion(() ->{
            log.debug("SSE Connection closed");
            setHistory(conversationId, AIConstant.ROLE_ASSISTANT, responseBuilder.toString());
        });
        emitter.onTimeout(() -> {
            log.debug("SSE Timeout");
        });
        emitter.onError(e -> {
            log.debug("SSE Error: " + e.getMessage());
        });

        return emitter;
    }

    private void setHistory(Long conversationId, String role, String prompt) {
        ChatQueues<Map<String, String>> queues = ChatQueues.getInstance();
        FixedSizeQueue<Map<String, String>> historyMsgQueue = queues.getQueue(conversationId);
        Map<String, String> msgMap = new HashMap<>(1);
        msgMap.put("role", role);
        msgMap.put("content", prompt);
        historyMsgQueue.offer(msgMap);
    }

    private List<Map<String, String>> getHistory(Long conversationId) {
        ChatQueues<Map<String, String>> queues = ChatQueues.getInstance();
        FixedSizeQueue<Map<String, String>> historyMsgQueue = queues.getQueue(conversationId);
        List<Map<String,String>> result = new ArrayList<>();
        while(!historyMsgQueue.isEmpty()) {
            Map<String, String> poll = historyMsgQueue.poll();
            result.add(poll);
        }
        return result;
    }


    private void buildBody(AiChatForm form, Map<String, Object> requestBody){

        List<Map<String,String>> promptMessages = new ArrayList<>();
        //查询历史会话
        List<Map<String,String>> historyMessages = getHistory(form.getConversationId());

        if(!CollectionUtils.isEmpty(historyMessages)){
            promptMessages.addAll(historyMessages);
        }
        promptMessages.add(
                new HashMap<String, String>() {{
                    put("content", form.getPrompt());
                    put("role", AIConstant.ROLE_USER);
                }});

        if("1".equals(form.getThinkingEnabled())) {
            requestBody.put("model", aiConfig.getThinkmodel());
        } else {
            requestBody.put("model", aiConfig.getModel());
        }
        requestBody.put("messages", promptMessages);
        if (aiConfig.isToolsEnable()) {
            requestBody.put("tools", FunctionRegistry.getInstance().tools());
        }
        if(aiConfig.getMaxTokens()!=null){
            requestBody.put("max_tokens", aiConfig.getMaxTokens());
        }else {
            requestBody.put("max_tokens", 4096);
        }
        requestBody.put("stream", true); // 启用流式模式
    }

    private int getTotalTokens(JsonNode usage){

        if (usage != null && usage.has("total_tokens")) {
            return usage.get("total_tokens").asInt();
        }
        return 0;
    }
}
